from distutils.log import info
from hashlib import new
from re import sub
import rdkit
import os
import numpy as np
import pandas as pd
from biopandas.pdb import PandasPdb
import glob
import argparse
import subprocess

parser = argparse.ArgumentParser(description="pdb file to be changed")

parser.add_argument('-pdb_in', type=str,
                    help="name of the pdb file i.e. (xxx.pdb)")
parser.add_argument('-pdb_out', type=str,
                    help="name of the modified pdb file i.e. (xxx.pdb)")
parser.add_argument('-res', type=str)
parser.add_argument('-atom_type_from', type=str)
parser.add_argument('-atom_type_to', type=str)
parser.add_argument('-C_end_O1',type=str)
parser.add_argument('-C_end_O2',type=str)
parser.add_argument('-C_end_resn',type=str)

args = parser.parse_args()
pdb_fname_in=args.pdb_in
pdb_fname_out=args.pdb_out
res_modify=args.res
atom_type_i=args.atom_type_from
atom_type_f=args.atom_type_to
c_term_o1_type=args.C_end_O1
c_term_o2_type=args.C_end_O2
c_term_resn=args.C_end_resn

data=PandasPdb().read_pdb(pdb_fname_in)
df=data.df['ATOM']
tmp=df[(df.residue_name==res_modify) & (df.atom_name==atom_type_i)]
tmp1=tmp.replace(atom_type_i,atom_type_f)
replace_idx=tmp1.index.values
df.iloc[replace_idx,:]=tmp1
data.df['ATOM']=df

df2=data.df['ATOM']
tmp2=df2[df.residue_name==c_term_resn]
tmp3=tmp2.iloc[-2:]
tmp4=tmp3.replace(tmp3.iloc[-2,3],c_term_o1_type)
tmp5=tmp4.replace(tmp3.iloc[-1,3],c_term_o2_type)
replace_idx=tmp5.index.values
df2.iloc[replace_idx,:]=tmp5
data.df['ATOM']=df2
data.df['ATOM']=data.df['ATOM'][data.df['ATOM']['element_symbol'] != 'H']
data.to_pdb(pdb_fname_out)